import time
import tiktoken
from functools import wraps
import logging


##拓展装饰器
##统计时间消耗
def timecost(func):
    # 使用 functools.wraps 来保留原始函数的元数据
    @wraps(func)
    def wrapper(*args, **kwargs):
        # 获取当前日期和时间
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        end_time = time.perf_counter()
        ##计算执行消耗时间
        print(f"函数 {func.__name__} 的执行时间为: {end_time - start_time} 秒")
        logging.info(f"函数 {func.__name__} 的执行时间为: {end_time - start_time} 秒")
        return result
    return wrapper

##统计token消耗
def tokencost(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        if len(result) == 0:
            return True
            # 初始化 GPT-3 编码器
        encoding = tiktoken.get_encoding("cl100k_base")
        # 计算每个字符串的 token 数量
        token_counts_list = [len(encoding.encode(string)) for string in result]
        total_consume = sum(token_counts_list)
        print("历史对话记录消耗token为"+str(total_consume))
        return result
    return wrapper

